#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>

int main(int argc, char** argv) {
    int  rank, size;
    int* sendbuf;
    int* sendcounts;
    int* displs;
    int  recvbuf[3]; // 假设每个进程最多接收 3 个元素

    // 初始化 MPI 环境
    MPI_Init(&argc, &argv);

    // 获取当前进程的 rank 和总的进程数
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    // 只有根进程（rank 0）初始化发送缓冲区
    if (rank == 0) {
        sendbuf = (int*)malloc(10 * sizeof(int)); // 假设总共发送 10 个元素
        for (int i = 0; i < 10; i++) {
            sendbuf[i] = i + 1;
        }

        // 每个进程接收的元素数量
        sendcounts = (int*)malloc(size * sizeof(int));
        sendcounts[0] = 2;
        sendcounts[1] = 3;
        sendcounts[2] = 2;
        sendcounts[3] = 3;

        // 每个进程接收数据的偏移量
        displs = (int*)malloc(size * sizeof(int));
        displs[0] = 0;
        displs[1] = 2;
        displs[2] = 5;
        displs[3] = 7;
    }

    // 使用 MPI_Scatterv 分发数据
    // sendbuf：   发送缓冲区的起始地址（仅对根进程有意义）。
    // sendcounts：一个整数数组，表示每个进程接收的元素数量（仅对根进程有意义）。
    // displs：    一个整数数组，表示每个进程接收数据在发送缓冲区中的偏移量（仅对根进程有意义）。
    // sendtype：  发送缓冲区中元素的数据类型。
    // recvbuf：   接收缓冲区的起始地址。
    // recvcount： 接收缓冲区中元素的数量。
    // recvtype：  接收缓冲区中元素的数据类型。
    // root：      根进程的 rank。
    // comm：      通信子。
    MPI_Scatterv(sendbuf, sendcounts, displs, MPI_INT, recvbuf, 3, MPI_INT, 0, MPI_COMM_WORLD);

    // 每个进程打印接收到的数据
    printf("Process %d received data: ", rank);
    for (int i = 0; i < sendcounts[rank]; i++) {
        printf("%d ", recvbuf[i]);
    }
    printf("\n");

    // 根进程释放发送缓冲区
    if (rank == 0) {
        free(sendbuf);
        free(sendcounts);
        free(displs);
    }

    // 结束 MPI 环境
    MPI_Finalize();

    return 0;
}